package sample.methodframe;

import com.google.common.collect.ImmutableMap;
import org.apache.commons.collections4.CollectionUtils;
import org.objectweb.asm.*;
import org.objectweb.asm.commons.AnalyzerAdapter;
import util.FileUtils;

import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.stream.Collectors;

public class MethodFramePrintDemo {
    public static void main(String[] args) throws IOException {
        byte[] bytes = invoke(HelloWorld.class);
        //  output
        FileUtils.writeBytes(Type.getInternalName(HelloWorld.class) + "MethodAround.class", bytes);
        //  replace
        FileUtils.writeBytes(HelloWorld.class.getResource("/").getPath() + Type.getInternalName(HelloWorld.class) + ".class", bytes);
    }

    private static byte[] invoke(Class<?> clazz) throws IOException {
        ClassReader classReader = new ClassReader(clazz.getName());
        ClassWriter classWriter = new ClassWriter(classReader, ClassWriter.COMPUTE_FRAMES);
        MethodFramePrintVisitor visitor = new MethodFramePrintVisitor(Opcodes.ASM9, classWriter);
        classReader.accept(visitor, ClassReader.EXPAND_FRAMES);
        return classWriter.toByteArray();
    }

    protected static class MethodFramePrintVisitor extends ClassVisitor {
        private String owner;

        public MethodFramePrintVisitor(int api, ClassVisitor classVisitor) {
            super(api, classVisitor);
        }

        @Override
        public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {
            super.visit(version, access, name, signature, superName, interfaces);
            this.owner = name;
        }

        @Override
        public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
            MethodVisitor mv = super.visitMethod(access, name, descriptor, signature, exceptions);
            return new MethodFramePrintAdapter(api, owner, access, name, descriptor, mv);
        }

        private class MethodFramePrintAdapter extends AnalyzerAdapter {
            public final Map<Integer, String> TYPE_MAP = new ImmutableMap.Builder<Integer, String>()
                    .put(Opcodes.TOP, "top")
                    .put(Opcodes.INTEGER, "int")
                    .put(Opcodes.FLOAT, "float")
                    .put(Opcodes.DOUBLE, "double")
                    .put(Opcodes.LONG, "long")
                    .put(Opcodes.NULL, "null")
                    .put(Opcodes.UNINITIALIZED_THIS, "uninitialized_this")
                    .build();
            private final Function<Object, String> labelFunction;
            private String methodName;
            private String methodDesc;

            protected MethodFramePrintAdapter(int api, String owner, int access, String name, String descriptor, MethodVisitor methodVisitor) {
                super(api, owner, access, name, descriptor, methodVisitor);
                this.methodName = name;
                this.methodDesc = descriptor;

                labelFunction = obj -> {
                    if (obj instanceof Label) {
                        Object uninitializedVal = uninitializedTypes.get(obj);
                        return uninitializedVal == null ? obj.toString() : "uninitialized_" + uninitializedVal;
                    }
                    return obj.toString();
                };
            }

            @Override
            public void visitCode() {
                super.visitCode();
                System.out.println();
                System.out.println(String.format("%s.%s:%s)", owner, methodName, methodDesc));
                printStackMapFrame();
            }

            @Override
            public void visitInsn(int opcode) {
                super.visitInsn(opcode);
                printStackMapFrame();
            }

            @Override
            public void visitIntInsn(int opcode, int operand) {
                super.visitIntInsn(opcode, operand);
                printStackMapFrame();
            }

            @Override
            public void visitVarInsn(int opcode, int var) {
                super.visitVarInsn(opcode, var);
                printStackMapFrame();
            }

            @Override
            public void visitTypeInsn(int opcode, String type) {
                super.visitTypeInsn(opcode, type);
                printStackMapFrame();
            }

            @Override
            public void visitFieldInsn(int opcode, String owner, String name, String descriptor) {
                super.visitFieldInsn(opcode, owner, name, descriptor);
                printStackMapFrame();
            }

            @Override
            public void visitMethodInsn(int opcodeAndSource, String owner, String name, String descriptor, boolean isInterface) {
                super.visitMethodInsn(opcodeAndSource, owner, name, descriptor, isInterface);
                printStackMapFrame();
            }

            @Override
            public void visitInvokeDynamicInsn(String name, String descriptor, Handle bootstrapMethodHandle, Object... bootstrapMethodArguments) {
                super.visitInvokeDynamicInsn(name, descriptor, bootstrapMethodHandle, bootstrapMethodArguments);
                printStackMapFrame();
            }

            @Override
            public void visitJumpInsn(int opcode, Label label) {
                super.visitJumpInsn(opcode, label);
                printStackMapFrame();
            }

            @Override
            public void visitLdcInsn(Object value) {
                super.visitLdcInsn(value);
                printStackMapFrame();
            }

            @Override
            public void visitIincInsn(int var, int increment) {
                super.visitIincInsn(var, increment);
                printStackMapFrame();
            }

            @Override
            public void visitTableSwitchInsn(int min, int max, Label dflt, Label... labels) {
                super.visitTableSwitchInsn(min, max, dflt, labels);
                printStackMapFrame();
            }

            @Override
            public void visitLookupSwitchInsn(Label dflt, int[] keys, Label[] labels) {
                super.visitLookupSwitchInsn(dflt, keys, labels);
                printStackMapFrame();
            }

            @Override
            public void visitMultiANewArrayInsn(String descriptor, int numDimensions) {
                super.visitMultiANewArrayInsn(descriptor, numDimensions);
                printStackMapFrame();
            }

            @Override
            public void visitTryCatchBlock(Label start, Label end, Label handler, String type) {
                super.visitTryCatchBlock(start, end, handler, type);
                printStackMapFrame();
            }

            private void printStackMapFrame() {
                CollectionUtils.isEmpty(locals);
                String localsStr = list2Str(locals);
                String stackStr = list2Str(stack);
                System.out.println(String.format("locals: %s, stack: %s", localsStr, stackStr));
            }

            private String list2Str(List<Object> locals) {
                if (CollectionUtils.isEmpty(locals)) {
                    return "[]";
                }
                return locals.stream().map(this::item2Str).collect(Collectors.joining(" ", "[ ", " ]"));
            }

            private String item2Str(Object item) {
                return TYPE_MAP.getOrDefault(item, labelFunction.apply(item));
            }
        }
    }
}
